package com.example.homeworktagger.service;

import com.example.homeworktagger.model.Question;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.embedding.EmbeddingModel;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.SearchResults;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.collection.LoadCollectionParam;
import io.milvus.param.collection.ReleaseCollectionParam;
import io.milvus.param.dml.SearchParam;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;

/**
 * Milvus向量数据库服务类
 * 
 * 负责与Milvus数据库进行交互，实现向量相似度搜索功能
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class MilvusService {

    /** Milvus客户端 */
    private final MilvusServiceClient milvusClient;
    
    /** 向量嵌入模型，用于将文本转换为向量 */
    private final EmbeddingModel embeddingModel;

    /** 集合名称 */
    private static final String COLLECTION_NAME = "questions";
    
    /** 向量字段名称 */
    private static final String VECTOR_FIELD = "question_vector";
    
    /** 问题文本字段名称 */
    private static final String QUESTION_FIELD = "question";
    
    /** 答案文本字段名称 */
    private static final String ANSWER_FIELD = "answer";
    
    /** 检索结果数量限制 */
    private static final int TOP_K = 3;
    
    /** 相似度阈值，低于此值的结果将被过滤 */
    private static final float SIMILARITY_THRESHOLD = 0.5f;

    /**
     * 搜索相似题目
     * 
     * @param query 查询文本
     * @return 相似题目列表，按相似度降序排列
     */
    public List<Question> searchSimilarQuestions(String query) {
        try {
            // 1. 加载集合
            loadCollection();

            // 2. 生成查询文本的向量表示
            List<Float> queryVector = textToVector(query);

            // 3. 执行向量搜索
            SearchResults results = executeSearch(queryVector);

            // 4. 处理搜索结果
            return processSearchResults(results);
        } finally {
            // 5. 释放集合资源
            releaseCollection();
        }
    }

    /**
     * 加载Milvus集合到内存
     */
    private void loadCollection() {
        milvusClient.loadCollection(
            LoadCollectionParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .build()
        );
        log.debug("已加载集合: {}", COLLECTION_NAME);
    }

    /**
     * 将文本转换为向量
     * 
     * @param text 输入文本
     * @return 向量表示（浮点数列表）
     */
    private List<Float> textToVector(String text) {
        Embedding embedding = embeddingModel.embed(text).content();
        List<Float> vector = new ArrayList<>();
        for (float value : embedding.vector()) {
            vector.add(value);
        }
        return vector;
    }

    /**
     * 执行向量搜索
     * 
     * @param queryVector 查询向量
     * @return 搜索结果
     */
    private SearchResults executeSearch(List<Float> queryVector) {
        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .withVectorFieldName(VECTOR_FIELD)
                .withVectors(List.of(queryVector))
                .withTopK(TOP_K)
                .withMetricType(MetricType.COSINE)
                .withOutFields(List.of(QUESTION_FIELD, ANSWER_FIELD))
                .build();

        R<SearchResults> searchResults = milvusClient.search(searchParam);
        if (searchResults.getStatus() != R.Status.Success.getCode()) {
            throw new RuntimeException("搜索失败: " + searchResults.getMessage());
        }
        
        return searchResults.getData();
    }

    /**
     * 处理搜索结果，转换为Question对象列表
     * 
     * @param results 搜索结果
     * @return Question对象列表
     */
    private List<Question> processSearchResults(SearchResults results) {
        List<Question> similarQuestions = new ArrayList<>();
        
        for (int i = 0; i < results.getResults().getFieldsDataCount(); i++) {
            float similarity = results.getResults().getScores(i);
            
            // 过滤掉相似度低于阈值的结果
            if (similarity >= SIMILARITY_THRESHOLD) {
                Question question = new Question();
                question.setQuestion(results.getResults().getFieldsData(0).getScalars().getStringData().getData(i));
                question.setAnswer(results.getResults().getFieldsData(1).getScalars().getStringData().getData(i));
                question.setSimilarity(similarity);
                similarQuestions.add(question);
                
                log.info("找到相似题目：{}，相似度：{}", question.getQuestion(), similarity);
            }
        }
        
        return similarQuestions;
    }

    /**
     * 释放集合资源
     */
    private void releaseCollection() {
        milvusClient.releaseCollection(
            ReleaseCollectionParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .build()
        );
        log.debug("已释放集合: {}", COLLECTION_NAME);
    }
} 